package cn.kennylee.codehub.springboot.microtrace;

import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.lang.Nullable;
import org.springframework.util.StringUtils;
import org.springframework.web.server.ServerWebExchange;

import static cn.kennylee.codehub.springboot.microtrace.TraceContext.generateTraceId;


/**
 * <p> 微服务请求跟踪工具类 </p>
 * <p>Created on 2025/5/10.</p>
 *
 * @author kennylee
 * @since 0.0.1
 */
@Slf4j
public final class TraceUtils {

    /**
     * 判断MDC是否已经包含Trace ID
     *
     * @return true: 包含Trace ID; false: 不包含Trace ID
     */
    public static boolean hasTraceId() {
        return StringUtils.hasLength(MDC.get(getMdcTraceIdName()));
    }


    /**
     * 获取Trace ID在header中的key名称
     *
     * @return Trace ID在header中的key名称
     */
    public static String getHeaderTraceIdName() {
        return SpringUtil.getBean(TraceProperties.class).getTraceHeaderName();
    }

    /**
     * 获取Trace ID在MDC中的key名称
     *
     * @return Trace ID在MDC中的key名称
     */
    public static String getMdcTraceIdName() {
        return SpringUtil.getBean(TraceProperties.class).getMdcTraceIdName();
    }

    /**
     * 处理是否接收传递的 Trace ID
     *
     * @param transferTraceId 传递的Trace ID
     * @return Trace ID
     */
    @Nullable
    public static String acceptTransferTraceId(@Nullable String transferTraceId) {

        TraceProperties traceProperties = SpringUtil.getBean(TraceProperties.class);

        String traceId = transferTraceId;
        // 不接受传递的 Trace ID
        if (!traceProperties.isAcceptTransferTraceId()) {
            String msg = "";
            if (StringUtils.hasLength(transferTraceId)) {
                msg += "不接收传递的";
            } else {
                msg += "不存在";
            }
            traceId = generateTraceId();
            log.debug("{} Trace ID，新生成: {}", msg, traceId);
        }
        return traceId;
    }

    /**
     * 获取Trace ID
     *
     * @param serverWebExchange ServerWebExchange
     * @return Trace ID
     */
    @Nullable
    public static String getTraceId(ServerWebExchange serverWebExchange) {
        String traceId = serverWebExchange.getRequest()
            .getHeaders()
            .getFirst(getHeaderTraceIdName());

        // 如果请求头中没有 Trace ID，尝试从查询参数
        TraceProperties traceProperties = SpringUtil.getBean(TraceProperties.class);
        if (!StringUtils.hasLength(traceId) && traceProperties.isAcceptQueryTraceId()) {
            traceId = StrUtil.trimToNull(serverWebExchange.getRequest().getQueryParams().getFirst(traceProperties.getMdcTraceIdName()));
        }

        return traceId;
    }
}
